#ifndef CVX_ADAPTIVE_QUAD_TREE_H
#define CVX_ADAPTIVE_QUAD_TREE_H

#include "cvxImage.h"
#include "cvxQuadTree.h"

/************************************************************************/
/* 
	adaptive quad tree according to image size
	children is NULL when pixel position isn't in image scope
*/
/************************************************************************/
class CvxAdaptiveQuadTree: public CvxQuadTree
{
public:
	CvxAdaptiveQuadTree(unsigned int level, const Mat &img, int threshold = 5);
	CvxAdaptiveQuadTree(const CvxAdaptiveQuadTree &tree);
	~CvxAdaptiveQuadTree();

	void create(QtNode * node);
	void balance();
	void reverse(QtNode *node);
	int unknownNodeNumber() {return m_unknownNum; }
	void updateUnknownIndex(const int baseIndex, QtNode *node);  //update unkonwn node index, e.g right image
	void matchNode(float rgbThreshold, const int sizeThreshold, QtNode *node, CvxAdaptiveQuadTree *otherTree);    //match leaf node according to color distribution and size
	void crossMatchNode(QtNode *node, CvxAdaptiveQuadTree *otherTree);

	//initLeafData data in tree leaf node, idx, shading and average color
	void initLeafData(const Mat &img, const Mat &shadingMap, const Mat &knownMap, 
						QtNode *node, const Mat&edgeMap);

	void adaptiveInitLeafData(const Mat &img, const Mat &edgMap, QtNode *node, int threshold = 5, float minimumShadingScale = 1.0);
	
	//cluster larger node to similar specular
	void clusterLargeNode(QtNode *node, const int sizeThreshold, const Mat &labelMap, const int labelNum);

	//update edge node shading by nearby node
	void updateEdgeNodeShading(QtNode *node, Mat &shadingMap,  Mat &edgeMap, Mat &albedo, int winSize = 5);

	QtNode *LocateNode (int x, int y);
	QtNode *LocateRightNeighbor(QtNode *cell);
	QtNode *LocateTopNeighbor(QtNode *cell);

	//after balance, locate neighbors
	//0, 1 or 2 neighbors 
	int leftNeighbors(QtNode *cell, QtNode *&n1, QtNode *&n2);
	int bottomNeighbors(QtNode *cell, QtNode *&n1, QtNode *&n2);
	int rightNeighbors(QtNode *cell, QtNode *&n1, QtNode *&n2);
	int topNeighbors(QtNode *cell, QtNode *&n1, QtNode *&n2);

	void ConstructRestrictedQuadTree(QtNode *node, bool& changed);

	//test
	void draw(QtNode * node);
	void drawColorNode(QtNode *node, const Scalar& colorA, const Scalar& colorB);         //solved and unsolved node has different color
	void drawKnownShading(QtNode *node, Mat &shadingMap);
	void drawMatchNode(QtNode *node, CvxAdaptiveQuadTree *otherTree, Mat &resultImage);                    // draw corresponding node
	void compareGroundTruthWithGrayLevelInPixelLeaf(QtNode *node, Mat &shadingGroundTruth, vector<unsigned int> &positionIdx);

	//   draw clustered node which similar color
	void drawClusteredNode(QtNode *node, int labelNum, vector<Scalar> &colors);
	
	//test only
	void setShadingToEdgeNode(QtNode *node,  Mat &edgeMap, const Mat &shadingMap);

	//bilinear interpolate shading from node
	//void InterpolateShading(QtNode *node, Mat &shadingMap);
	//interpolate a vector from node
	void InterpolateAvector(QtNode *node, Mat &AvectorMap);

	

	Mat m_orgimg;
	Mat m_orgBackUp;                   //back up image of the original image for visualization purpose
private:
	int m_threshold;                   //Std Dev of similar color definition
	int m_initKnownNodeThreshold;      //threshold that set as initial known node     
	int m_unknownNum;
	vector<Mat> m_ShadingMatVec;       //matrix vector to speed up interpolate by avoid Mat create, 
	vector<Mat> m_AvectorMatVec;       //matrix vector for A vector interpolation

public:

	
	//set shading value to the node and shading image
	void setShadingResult(Mat &shaImg, QtNode *node, const vector<double> &rst);
	void upsampleShadingResult(Mat &shaImg, QtNode *node);

	//equation 10 in the paper "User-assisted intrinsic images"
	void planeEquation(QtNode *node, vector<map<int, double>> &leftMap, vector<double> &rightVec, bool appliedToEdgeArea = false);
	//set equation on edge pixel areas
	void edgeEquation(QtNode *node, vector<map<int, double>> &leftMap, vector<double> &rightVec);
	
	//set equation on matched nodes: coherent reflectance
	void matchNodeEquation(QtNode *node, vector<map<int, double>> &leftMap, vector<double> &rightVec);

	//set equation according equ 8 to get a_3_1 vector
	void AvectorEquation(QtNode *node, vector<map<int, double>> &leftMap, vector<double> &rightVec);
	
	//set A vector to an image
	void setAvector(Mat &AvectorMap, QtNode *node, const vector<double> &rst);

	//set A vector to node
	void setAvector(QtNode *node, const vector<double> &rst);

};

#endif